#include "socket.h"
#include "log.h"

#ifdef _WIN32
typedef int socklen_t;
#include <time.h>
#else
#define INVALID_SOCKET (~0)
#define SOCKET_ERROR (-1)
#define SD_SEND (1)
#define closesocket(x) close(x)
#include <sys/time.h>
#include <string.h>
#include <arpa/inet.h>
#include <fcntl.h>
#endif

static bool WSAStartupCalled = false;

int GetFileSize(const char* file)
{
  FILE* fp = fopen(file, "rb");
  fseek(fp, 0, SEEK_END);
  int len = ftell(fp);
  fclose(fp);
  return len;
}

Socket::Socket()
: connected(false)
, ssl(0)
, listener(INVALID_SOCKET)
, connection(0)
{
  memset(&connectionAddr, 0, sizeof(connectionAddr));

#ifdef _WIN32
  if (!WSAStartupCalled)
  {
    WSADATA data;
    WSAStartup(MAKEWORD(2,2), &data);
    WSAStartupCalled = true;
  }
#endif
}

Socket::~Socket()
{
  Close();
  delete ssl;
}

void Socket::Close()
{
  if (listener != INVALID_SOCKET)
    closesocket(listener);
}

void Socket::CloseConnection()
{
  if (connection != INVALID_SOCKET)
    closesocket(connection);
}

void Socket::EnableSSL()
{
  ssl = new SSL(); 
}

void Socket::Connect(const std::string& address_, int port)
{
  sockaddr_in address;
  memset(&address, 0, sizeof(address));
  address.sin_family = AF_INET;
  address.sin_addr.s_addr = inet_addr(address_.c_str());
  address.sin_port = htons((unsigned short)port);

  connection = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP);
  if (connection == INVALID_SOCKET) 
  {
    Log::Error("Socket::Listen creating socket failed, port %d.", port);
    return;
  }

  if (connect(connection, (sockaddr*)&address, sizeof(address)) == SOCKET_ERROR) 
  {
    Log::Error("Socket::Connect failed [%s:%d]", address_.c_str(), port);
    return;
  }
  connected = true;
}

bool Socket::Listen(int port)
{
  sockaddr_in address;
  memset(&address, 0, sizeof(address));
  address.sin_family = AF_INET;
  address.sin_port = htons((unsigned short)port);

  listener = socket(PF_INET, SOCK_STREAM, 6);
  if (listener == INVALID_SOCKET) 
  {
    Log::Error("Socket::Listen creating socket failed, port %d.", port);
    return false;
  }
  int optionValue = 1;
  if (setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, (char*)&optionValue, sizeof(optionValue)) != 0)
  {
    Log::Error("Socket::Listen setsockopt() failed.");
    closesocket(listener);
    return false;
  }
  if (bind(listener, (struct sockaddr *)&address, sizeof(address)) != 0)
  {
    Log::Error("Socket::Listen bind() failed.");
    closesocket(listener);
    return false;
  }
  if (listen(listener, 20) != 0)
  {
    Log::Error("Socket::Listen listen() failed.");
    closesocket(listener);
    return false;
  }

  #ifndef _WIN32
    fcntl(listener, F_SETFD, FD_CLOEXEC);
  #endif
  return true;
}

void Socket::WriteEnd()
{
  if (!connected)
    return;
  int result = shutdown(connection, SD_SEND);
  Log::ErrorIf(result == SOCKET_ERROR, "Socket::Shutdown failed");
  if (result == SOCKET_ERROR) 
    closesocket(connection);
}

int Socket::CheckConnections(const std::vector<Socket*>& sockets, int timeOutMillisec)
{
  fd_set read_set;
  FD_ZERO(&read_set);
  int max_fd = -1;
  for (size_t i = 0; i < sockets.size(); ++i)
  {
    FD_SET(sockets[i]->listener, &read_set);
    if ((int)sockets[i]->listener > max_fd)
      max_fd = (int)sockets[i]->listener;
  }

  struct timeval timeout;
  timeout.tv_sec = (timeOutMillisec / 1000);
  timeout.tv_usec = (timeOutMillisec % 1000) * 1000;
  int numConnection = select(max_fd + 1, &read_set, 0, 0, &timeout);
#ifdef _WIN32
  int error = WSAGetLastError();
#endif
  if (numConnection >= 0)
  {
    for (size_t i = 0; i < sockets.size(); ++i)
      if (FD_ISSET(sockets[i]->listener, &read_set))
        return (int)i;
  }
  return -1;
}

void Socket::AcceptConnection()
{
  socklen_t length = sizeof(connectionAddr);
  connection = accept(listener, &connectionAddr, &length);
  Log::ErrorIf(connection == INVALID_SOCKET, "Socket::AcceptConnection accept failed.");
  if (connection == INVALID_SOCKET) 
    return;
  if (ssl != 0)
    ssl->SetSocket((int)connection);
  connected = true;
}

std::string Socket::ReadString(int timeOutMillisec)
{
  std::string ret;
  time_t start = time(0);
  double timeoutSeconds = ((double)timeOutMillisec) / 1000.0;
  double elapsedSeconds = 0;
  while (elapsedSeconds < timeoutSeconds && ret.size() == 0)
  {
    ret += ReadString();
    time_t now = time(0);
    elapsedSeconds = difftime(now, start);
  }
  return ret;
}

std::string Socket::ReadString()
{
  if (!connected)
    return std::string();

  int read = 0;
  char buf[8192];

  if (ssl != 0)
    read = ssl->Read(buf, sizeof(buf) - 1);
  else 
    read = recv(connection, buf, sizeof(buf) - 1, 0);
  
  if (read == -1)
    return std::string();
  buf[read] = 0;
  return buf;
}

void Socket::WriteString(const std::string& s)
{
  if (!connected)
    return;
  if (ssl != 0)
    ssl->Write(&s[0], (int)s.size());
  else
    send(connection, &s[0], (int)s.size(), 0);
}

void Socket::WriteBinary(const std::string& file, int range1, int range2)
{
  if (!connected)
    return;
  enum { buffSize = 32768 };
  if (range2 == 0)
    range2 = GetFileSize(file.c_str());
  int length = range2 - range1 + 1;
  unsigned char buff[buffSize];
  FILE* fp = fopen(file.c_str(), "rb");
  Log::ErrorIf(fp == 0, "Socket::WriteBinary file not found '%s'", file.c_str());
  if (fp == 0)
    return;
  fseek(fp, range1, SEEK_SET);
  while (length > 0)
  {
    int read = (buffSize > length ? length : buffSize);
    fread(buff, 1, read, fp);
    if (ssl != 0)
      ssl->Write((const char*)&buff, read);
    else
      send(connection, (const char*)&buff, read, 0);
    length -= read;
  }
  fclose(fp);
}
